
import librosa
import argparse
import numpy as np
import librosa.display
# import SpecAugment
import matplotlib.pyplot as plt
import random
from tensorflow_addons.image import sparse_image_warp


def getMel(filepath, name, makirpath): #获取频谱图
    y, sr = librosa.load(path=filepath)
    melspec = librosa.feature.melspectrogram(y=y, sr=sr, hop_length=128, n_mels=256, fmax=8000)
    logmelspec = librosa.power_to_db(melspec)
    # plt.figure()
    fig = plt.figure(figsize=(12.8, 12.8), dpi=20, frameon=False)  # 像素点20，不显示边框
    heatmap = plt.pcolor(logmelspec)  # 根据spec的值，画色彩图
    plt.xticks([])
    plt.yticks([])
    plt.tight_layout()
    path = makirpath+'\\'+name+'.png'

    plt.savefig(path)
    plt.clf()





def time_warp(file, name, makirpath):#时间扭曲
    audio, sr = librosa.load(file)
    # Extract Mel Spectrogram Features from the audio file
    mel_spectrogram = librosa.feature.melspectrogram(y=audio, sr=sr, n_mels=256, hop_length=128, fmax=8000)

    mel_spectrogram = np.reshape(mel_spectrogram,
                                 (-1, mel_spectrogram.shape[0], mel_spectrogram.shape[1], 1))

    v, tau = mel_spectrogram.shape[1], mel_spectrogram.shape[2]

    horiz_line_thru_ctr = mel_spectrogram[0][v // 2]

    random_pt = horiz_line_thru_ctr[
        random.randrange(80, tau - 80)]  # random point along the horizontal/time axis
    w = np.random.uniform((-80), 80)  # distance

    # Source Points
    src_points = [[[v // 2, random_pt[0]]]]

    # Destination Points
    dest_points = [[[v // 2, random_pt[0] + w]]]
    mel_spectrogram, _ = sparse_image_warp(mel_spectrogram, src_points, dest_points,
                                           num_boundary_points=2)
    logmelspec = librosa.power_to_db(mel_spectrogram[0, :, :, 0], ref=np.max)
    plt.figure(figsize=(12.8, 12.8), dpi=20, frameon=False)  # 像素点20，不显示边框
    plt.pcolor(logmelspec)  # 根据spec的值，画色彩图
    plt.xticks([])
    plt.yticks([])
    plt.tight_layout()
    path = makirpath + '\\' + name + 'time_warp.png'
    plt.savefig(path)
    plt.clf()

def freq_mask(file, name, makirpath):#频率掩盖
    audio, sr = librosa.load(file)
    # Extract Mel Spectrogram Features from the audio file
    mel_spectrogram = librosa.feature.melspectrogram(y=audio, sr=sr, n_mels=256, hop_length=128, fmax=8000)
    v = mel_spectrogram.shape[1]  # no. of mel bins

    # apply m_F frequency masks to the mel spectrogram
    for i in range(1):
        f = int(np.random.uniform(0, 27))  # [0, F)
        f0 = random.randint(0, v - f)  # [0, v - f)
        print(mel_spectrogram[:, f0:f0 + f, :, :])
        # mel_spectrogram[:, f0:f0 + f, :, :] = 0
    logmelspec = librosa.power_to_db(mel_spectrogram[0, :, :, 0], ref=np.max)
    plt.figure(figsize=(12.8, 12.8), dpi=20, frameon=False)  # 像素点20，不显示边框
    plt.pcolor(logmelspec)  # 根据spec的值，画色彩图
    plt.xticks([])
    plt.yticks([])
    plt.tight_layout()
    path = makirpath + '\\' + name + 'freq_mask.png'
    plt.savefig(path)
    plt.clf()





